#####################################################################
# Figure 3
#####################################################################

#####################################################################
# Preliminaries
#####################################################################

rm(list = ls())

library(scinference)
library(limSolve)

set.seed(12345)

#####################################################################
# Functions
#####################################################################

source("code/auxiliary scripts/common_functions.R")

#####################################################################
# Calibration DGPs
#####################################################################

source("code/auxiliary scripts/calibration_dgps.R")

#####################################################################
# Simulations for small T0
#####################################################################

# Setup
nreps_tradeoff <- 20000

Ks  <- 2:5
T0s <- c(30,45,60,75,90)

# Simulations

overall_cov_Ks <- overall_leng_Ks <- matrix(NA,length(Ks),length(T0s))

for (k in 1:length(Ks)){
  
  for (t in 1:length(T0s)){
    
    # print(c(Ks[k],T0s[t]))
    
    res_temp_cov <- res_temp_leng <- matrix(NA,nreps_tradeoff,1)
    
    for (r in 1:nreps_tradeoff){
      
      res_temp  <- sim_one_sample(DGP=1,T0=T0s[t],T1=T1_app,J=J_app,K=Ks[k],Lambda,rho_u=0.6,var_u,var_factors,rho_vec,var_epsl_vec,w0_sc,tau_alternative=0,li=FALSE,sdid=FALSE,permtest=FALSE)

      res_temp_cov[r]   <- res_temp$cov_all[1]
      res_temp_leng[r]  <- res_temp$leng_all[1]
      
    }
    
    overall_cov_Ks[k,t]   <- colMeans(res_temp_cov)
    overall_leng_Ks[k,t]  <- colMeans(res_temp_leng)
    
  }
  
}

#####################################################################
# Figures 
#####################################################################

# Coloring is based on the examples here: https://rdrr.io/r/graphics/persp.html)

cov_axis <- c(0.7,1)
leng_axis <- c(0,0.5)

nrowz <- nrow(overall_cov_Ks)
ncolz <- ncol(overall_cov_Ks)
jet.colors <- colorRampPalette( c("orange", "blue") )
ncol <- 100
color <- jet.colors(ncol)

graphics.off()
midpoints_cov <- overall_cov_Ks[-1, -1] + overall_cov_Ks[-1, -ncolz] + overall_cov_Ks[-nrowz, -1] + overall_cov_Ks[-nrowz, -ncolz]
pdf("graphics/3d_cov.pdf",pointsize=18,width=10.0,height=8.0)
pmat <- persp(x=Ks,y=T0s,z=overall_cov_Ks,col = color[cut(midpoints_cov,ncol)], xlim=range(Ks),ylim=range(T0s),
              zlim=range(cov_axis),xlab="K",ylab="Number of pre-treatment periods",zlab="Coverage",ticktype="detailed",expand=0.8,theta = 50,phi=20,d=10,nticks=4,main="Coverage 90% confidence intervals")
dev.off()

graphics.off()
midpoints_leng <- overall_leng_Ks[-1, -1] + overall_leng_Ks[-1, -ncolz] + overall_leng_Ks[-nrowz, -1] + overall_leng_Ks[-nrowz, -ncolz]
pdf("graphics/3d_leng.pdf",pointsize=18,width=10.0,height=8.0)
pmat <- persp(x=Ks,y=T0s,z=overall_leng_Ks,col = color[cut(midpoints_leng, ncol)],xlim=range(Ks),ylim=range(T0s),
              zlim=range(leng_axis),xlab="K",ylab="Number of pre-treatment periods",zlab="Average length",ticktype="detailed",expand=0.8,theta = 50,phi=20,d=10,nticks=4,main="Average length 90% confidence intervals")
dev.off()


